from langchain_chroma import Chroma

from langgraph.graph import MessagesState, START

from langchain_core.tools import create_retriever_tool
from langgraph.prebuilt import ToolNode
from langgraph.types import interrupt, Command
from pydantic import BaseModel
from langgraph.graph import END, StateGraph
from langgraph.checkpoint.memory import MemorySaver

from qw_model import qw_model

from Embed_model import BGE_Embed

# -------------------------------------定义工具
# 加载知识库
embedding = BGE_Embed()
vectorstore = Chroma(
    embedding_function=embedding,
    persist_directory='./chroma_db'
)

# 创建检索器
retriever = vectorstore.as_retriever(search_kwargs={'k': 3})

# 创建retriever工具
retriever_tool = create_retriever_tool(
    retriever,
    "rag",
    "对用户的业务咨询或者价格问题进行查询"
)

# 工具列表
tools = [retriever_tool]


# 定义一个模拟工具来询问用户问题
class AskHuman(BaseModel):
    """问用户一个问题"""
    question: str


# 创建模型，并绑定工具
model = qw_model()
model = model.bind_tools(tools + [AskHuman])


# 定义决定是否继续的函数
def should_continue(state):
    messages = state["messages"]
    last_message = messages[-1]

    # 如果没有函数调用，流程结束
    if not last_message.tool_calls:
        return END

    # 如果工具调用正在询问人类，返回"ask_human"节点
    elif last_message.tool_calls[0]["name"] == "AskHuman":
        return "ask_human"

    # 否则继续执行“action”节点
    else:
        return "action"


# 定义调用模型的函数
def call_model(state):
    messages = state["messages"]

    # 获取用户的问题
    user_message = messages[-1].content

    # 从检索器中获取相关文档
    context = retriever.invoke(user_message)

    # 格式化提示词
    prompt_template = """
    你是一个聪明又有亲和力的销售客服，名叫银弓，擅长用轻松、亲切的语气解答用户关于公司业务、价格、产品咨询的问题。
    这是检索到的文档: 
    {context}

    这是用户问题: {question}

    1. 当用户问题不清楚时，主动询问用户具体需求。
    2. 回答价格问题，严格按照检索的文档给出，不可自行编造。
    3. 当没有用户所说的地区，就按它的省份价格算，当省份也没有那就只能做不限地区的。
    4. 保持沟通语气亲切、自然，避免过于正式的措辞，多用口语化的表达和轻松的语气。
    5. 输出答案应简洁明了，字数在15字左右。
    """

    # 格式化提示词
    prompt = prompt_template.format(context=context, question=user_message)

    # 调用模型生成响应
    response = model.invoke([{"role": "user", "content": prompt}])

    return {"messages": [response]}


# 定义询问人类的函数
def ask_human(state):
    tool_call_id = state["messages"][-1].tool_calls[0]["id"]
    location = interrupt("请提供你的位置：")
    tool_message = [{"tool_call_id": tool_call_id, "type": "tool", "content": location}]
    return {"messages": tool_message}


# 定义状态图
workflow = StateGraph(MessagesState)

# 定义我们将在之间循环的三个节点
workflow.add_node("agent", call_model)
retrieve = ToolNode([retriever_tool])
workflow.add_node("action",retrieve)
workflow.add_node("ask_human", ask_human)

# 将入口点设置为 “代理”
workflow.add_edge(START, "agent")

# 定义条件边
workflow.add_conditional_edges(
    "agent",
    should_continue,
)

# 添加从“工具”到“代理”的边
workflow.add_edge("action", "agent")

# 在我们得到人类的反应后，我们回到agent那里
workflow.add_edge("ask_human", "agent")

# 内存和流处理
memory = MemorySaver()
app = workflow.compile(checkpointer=memory)
config = {"configurable": {"thread_id": "2"}}

# 模拟事件流
for event in app.stream(
        {
            "messages": [
                (
                        "user",
                        "使用检索工具询问用户需求，然后查找需求的价格或者产品咨询方面的回答",
                )
            ]
        },
        config,
        stream_mode="values",
):
    event["messages"][-1].pretty_print()

# 模拟另一个事件流
for event in app.stream(Command(resume="实探"), config, stream_mode="values"):
    event["messages"][-1].pretty_print()